from scipy.optimize import fsolve
from FunctionalBohm import SchrodingerNDT
from FunctionalBohm import Schrodinger1D
from math import pi

def segments(min,max,steps):
    return [i/float(steps)*(max-min)+min for i in range(steps+1)]

# Start solving:

from HarmonicOscillator1D import Psi as wvfn
from HarmonicOscillator1D import V
from HarmonicOscillator1D import mass

Schrod = SchrodingerNDT(V,wvfn,mass,vector_potential=lambda x:[0],dimensions=1,dx=0.001) # we need to specify a small enough stepsize!!!

# Alternatively we could have defined initial conditions of interest while defining Schrod, eg
# Schrod = SchrodingerNDT(..., positions=segments(-2,2,100), time=1.3, ...)
# But that's kind of boring because the positions are fixed, do not follow the flow (until I code that up later).

# get the physical variables of interest as functions of (t,x)

K = Schrod.get_Kinetic_Energy()
H = Schrod.get_Hamiltonian()
Q = Schrod.get_Quantum_Potential()

# We can now get Hamiltonian as a function of position (along rows) and time (along columns)
count = 0
for time in segments(0,4*pi/3.0,100): # A full period.#USING 4pi/3 ALSO: CREATE FILE FIRST!
    
    outfile = open("Z:/Desktop/Project/DATA/400data/400data" + "%03d" % (count) + ".csv",'w')
    count += 1

    for i in segments(-2,2,400):
        outfile.write(str(i)+", "+str(H(time,i))+", "+str(Q(time,i))+", "+str(V(i))+", "+str(K(time,i))+"\n")
    outfile.close()

print("Updated './data/'")



### To find level sets at height h, solve H - h == 0:
##
##def Levelset(H,t,h,start):
##    def diff(x): return H(t,x)-h
##    return fsolve(diff,start) # we need a vector of starting guesses...
##
##print(Levelset(H,t,6.25,[-0.75,-0.15,0.1,0.6]))
##
### NB: It would be nice to find the level sets without the starting guesses...
### Maybe just use segments(-2,2,400) as the guesses and then clean up duplicates later?
